import torch
import torchvision
from torch import nn
from torch.nn import Conv2d, MaxPool2d, Flatten, Sequential
from torch.utils.tensorboard import SummaryWriter
from torchvision import transforms

train_data = torchvision.datasets.CIFAR10(root='./dataset', train=True, download=True, transform=transforms.ToTensor())
train_loader = torch.utils.data.DataLoader(train_data, batch_size=128, shuffle=True)
class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.conv1 = Conv2d(in_channels=3, out_channels=32, kernel_size=5, stride=1, padding=2)
        self.maxPool = MaxPool2d(kernel_size=2, ceil_mode=True, padding=0)
        self.conv2 = Conv2d(in_channels=32, out_channels=32, kernel_size=5, stride=1, padding=2)
        self.conv3 = Conv2d(in_channels=32, out_channels=64, kernel_size=5, stride=1, padding=2)
        self.flatten = Flatten()
        self.linear1 = nn.Linear(in_features=64 * 4 * 4, out_features=64)
        self.linear2 = nn.Linear(in_features=64, out_features=10)

    def forward(self, x):
        x = self.conv1(x)
        x = self.maxPool(x)
        x = self.conv2(x)
        x = self.maxPool(x)
        x = self.conv3(x)
        x = self.maxPool(x)
        x = self.linear1(torch.flatten(x, 1))
        x = self.linear2(x)
        return x

model = MyModel()
print(model)


# Sequential用法

class ModelWithSequential(nn.Module):
    def __init__(self):
        super(ModelWithSequential, self).__init__()
        self.model1 = Sequential(
            Conv2d(in_channels=3, out_channels=32, kernel_size=5, stride=1, padding=2),
            MaxPool2d(kernel_size=2, ceil_mode=True, padding=0),
            Conv2d(in_channels=32, out_channels=32, kernel_size=5, stride=1, padding=2),
            MaxPool2d(kernel_size=2, ceil_mode=True, padding=0),
            Conv2d(in_channels=32, out_channels=64, kernel_size=5, stride=1, padding=2),
            MaxPool2d(kernel_size=2, ceil_mode=True, padding=0),
            Flatten(),
            nn.Linear(in_features=64 * 4 * 4, out_features=64),
            nn.Linear(in_features=64, out_features=10)
        )


    def forward(self, x):
        x = self.model1(x)
        return x

model1 = ModelWithSequential()
print(model1)

writer = SummaryWriter("./logs-seq")
step = 0

for data in train_loader:
    if step == 0:
        images, labels = data
        out_labels = model1(images)
        print(images.shape)
        writer.add_graph(model1, images)
    step += 1


writer.close()
